package com.gwz.comfyuiclient;

import cn.hutool.core.io.FileUtil;
import cn.hutool.poi.excel.ExcelReader;
import cn.hutool.poi.excel.ExcelUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;

import java.io.*;
import java.net.HttpURLConnection;
import java.net.URI;
import java.net.URL;
import java.net.URLEncoder;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.stream.Collectors;

public class LoadImageWebsocketComfyui {
    private static final String SERVER_ADDRESS = "127.0.0.1:8188";
    private final String clientId;
    private final ObjectMapper objectMapper;
    private Map<String, List<byte[]>> outputImages;
    private String currentNode;
    private static String BOOK_PATH = "C:\\Users\\shijian\\IdeaProjects\\comfyuiapi\\src\\main\\resources\\去除背景.json";

    public LoadImageWebsocketComfyui() {
        this.clientId = UUID.randomUUID().toString();
        this.objectMapper = new ObjectMapper();
        this.outputImages = new HashMap<>();
    }

    public String queuePrompt(Map<String, Object> prompt) throws IOException {
        Map<String, Object> data = new HashMap<>();
        data.put("prompt", prompt);
        data.put("client_id", clientId);

        URL url = new URL("http://" + SERVER_ADDRESS + "/prompt");
        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
        conn.setRequestMethod("POST");
        conn.setRequestProperty("Content-Type", "application/json");
        conn.setRequestProperty("Accept", "application/json");
        conn.setDoOutput(true);
        conn.setDoInput(true);

        String jsonInputString = objectMapper.writeValueAsString(data);
        try (OutputStream os = conn.getOutputStream()) {
            byte[] input = jsonInputString.getBytes(StandardCharsets.UTF_8);
            os.write(input, 0, input.length);
            os.flush();
        }

        if (conn.getResponseCode() != HttpURLConnection.HTTP_OK) {
            try (BufferedReader br = new BufferedReader(new InputStreamReader(conn.getErrorStream(), StandardCharsets.UTF_8))) {
                String errorBody = br.lines().collect(Collectors.joining("\n"));
                throw new IOException("HTTP错误代码: " + conn.getResponseCode() + "\n响应内容: " + errorBody);
            }
        }

        try (BufferedReader br = new BufferedReader(new InputStreamReader(conn.getInputStream(), StandardCharsets.UTF_8))) {
            Map<String, Object> response = objectMapper.readValue(br, Map.class);
            return (String) response.get("prompt_id");
        }
    }


    public byte[] getImage(String filename, String subfolder, String folderType) throws IOException {
        String query = String.format("filename=%s&subfolder=%s&type=%s",
                URLEncoder.encode(filename, StandardCharsets.UTF_8),
                URLEncoder.encode(subfolder, StandardCharsets.UTF_8),
                URLEncoder.encode(folderType, StandardCharsets.UTF_8));

        URL url = new URL("http://" + SERVER_ADDRESS + "/view?" + query);
        try (InputStream is = url.openStream()) {
            return is.readAllBytes();
        }
    }

    public Map<String, Object> getHistory(String promptId) throws IOException {
        URL url = new URL("http://" + SERVER_ADDRESS + "/history/" + promptId);
        try (BufferedReader br = new BufferedReader(new InputStreamReader(url.openStream()))) {
            return objectMapper.readValue(br, Map.class);
        }
    }

    public Map<String, List<byte[]>> getImages(Map<String, Object> prompt) throws Exception {
        outputImages.clear();
        String promptId = queuePrompt(prompt);

        WebSocketClient client = new WebSocketClient(new URI("ws://" + SERVER_ADDRESS + "/ws?clientId=" + clientId)) {
            @Override
            public void onMessage(String message) {
                try {
                    Map<String, Object> messageData = objectMapper.readValue(message, Map.class);
                    if ("executing".equals(messageData.get("type"))) {
                        Map<String, Object> data = (Map<String, Object>) messageData.get("data");
                        if (promptId.equals(data.get("prompt_id"))) {
                            currentNode = data.get("node") == null ? null : (String) data.get("node");
                        }
                    }
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }

            @Override
            public void onMessage(ByteBuffer bytes) {
                if ("save_image_websocket_node".equals(currentNode)) {
                    // 跳过前8个字节（与Python版本一致）
                    byte[] imageData = Arrays.copyOfRange(bytes.array(), 8, bytes.array().length);
                    outputImages.computeIfAbsent(currentNode, k -> new ArrayList<>()).add(imageData);
                }
            }

            @Override
            public void onOpen(ServerHandshake handshake) {
            }

            @Override
            public void onClose(int code, String reason, boolean remote) {
            }

            @Override
            public void onError(Exception ex) {
                ex.printStackTrace();
            }
        };

        client.connect();
        while (client.isOpen()) {
            if (currentNode == null) {
                break;
            }
            Thread.sleep(100);
        }
        client.close();

        return outputImages;
    }

    public static void main(String[] args) throws IOException {
        LoadImageWebsocketComfyui wsClient = new LoadImageWebsocketComfyui();

        for (File loopFile : FileUtil.loopFiles("C:\\Users\\shijian\\Desktop\\propicon")) {

            // 读取和修改提示配置
            String promptText = FileUtil.readString(BOOK_PATH, "UTF-8");


            Map<String, Object> prompt = wsClient.objectMapper.readValue(promptText, Map.class);
            ((Map<String, Object>) ((Map<String, Object>) prompt.get("2")).get("inputs")).put("image", loopFile.toPath().toString());

            ((Map<String, Object>) ((Map<String, Object>) prompt.get("4")).get("inputs")).put("filename_prefix", "gameIcon/" + loopFile.getName());
            wsClient.queuePrompt(prompt);
        }
    }
}